#!/usr/bin/env python3
"""
Complete client demo for Manga Colorizer API
Demonstrates full business workflow
"""

import requests
import time
import json
from PIL import Image
import io
import os
from typing import Dict, Any, Optional

API_BASE_URL = "http://localhost:8000/api/v1"

class MangaColorizerClient:
    """Complete client for Manga Colorizer API"""
    
    def __init__(self, base_url: str = API_BASE_URL):
        self.base_url = base_url
        self.session = requests.Session()
    
    def create_test_image(self, style: str = "manga") -> bytes:
        """Create a test manga sketch image"""
        img = Image.new('RGB', (512, 512), color='white')
        draw = Image.Draw(img)
        
        if style == "manga":
            # Manga character face
            draw.ellipse([(150, 100), (350, 300)], outline='black', width=3)
            draw.ellipse([(200, 150), (220, 170)], outline='black', width=2)  # Left eye
            draw.ellipse([(280, 150), (300, 170)], outline='black', width=2)  # Right eye
            draw.line([(200, 220), (300, 220)], fill='black', width=2)  # Mouth
            draw.arc([(140, 90), (360, 110)], 0, 180, fill='black', width=3)  # Hair
        elif style == "landscape":
            # Simple landscape
            draw.line([(0, 400), (512, 400)], fill='black', width=3)  # Horizon
            draw.ellipse([(200, 50), (300, 150)], outline='black', width=2)  # Sun
            draw.rectangle([(100, 300), (150, 400)], outline='black', width=2)  # House
        
        buffer = io.BytesIO()
        img.save(buffer, format='PNG', quality=95)
        return buffer.getvalue()
    
    def health_check(self) -> bool:
        """Check server health"""
        try:
            response = self.session.get(f"{self.base_url}/health", timeout=5)
            return response.status_code == 200 and response.json().get('status') == 'healthy'
        except Exception as e:
            print(f"❌ Health check failed: {e}")
            return False
    
    def upload_and_colorize(self, image_data: bytes, **kwargs) -> Optional[str]:
        """Upload image and start colorization"""
        files = {'file': ('manga_sketch.png', image_data, 'image/png')}
        data = {
            'style': kwargs.get('style', 'anime'),
            'strength': kwargs.get('strength', 0.7),
            'guidance_scale': kwargs.get('guidance_scale', 7.5),
            'num_inference_steps': kwargs.get('num_inference_steps', 20)
        }
        
        try:
            response = self.session.post(
                f"{self.base_url}/colorize",
                files=files,
                data=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                return result.get('task_id')
            else:
                print(f"❌ Upload failed: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            print(f"❌ Upload error: {e}")
            return None
    
    def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """Get task status"""
        try:
            response = self.session.get(
                f"{self.base_url}/colorize/status/{task_id}",
                timeout=10
            )
            return response.json() if response.status_code == 200 else {}
        except Exception as e:
            print(f"❌ Status check error: {e}")
            return {}
    
    def wait_for_completion(self, task_id: str, max_wait: int = 60) -> bool:
        """Wait for task completion"""
        print(f"⏳ Waiting for task {task_id} to complete...")
        
        for attempt in range(max_wait):
            status = self.get_task_status(task_id)
            current_status = status.get('status', 'unknown')
            progress = status.get('progress', 0)
            
            print(f"   Attempt {attempt + 1}: {current_status} ({progress}%)")
            
            if current_status == 'completed':
                print("✅ Task completed successfully!")
                return True
            elif current_status == 'failed':
                error_msg = status.get('error_message', 'Unknown error')
                print(f"❌ Task failed: {error_msg}")
                return False
            
            time.sleep(2)
        
        print("❌ Task timeout")
        return False
    
    def get_result(self, task_id: str) -> Optional[Dict[str, Any]]:
        """Get final result"""
        try:
            response = self.session.get(
                f"{self.base_url}/colorize/result/{task_id}",
                timeout=10
            )
            
            if response.status_code == 200:
                return response.json()
            else:
                print(f"❌ Result retrieval failed: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"❌ Result retrieval error: {e}")
            return None
    
    def batch_colorize(self, images: list, **kwargs) -> Optional[str]:
        """Batch colorize multiple images"""
        files = []
        for i, img_data in enumerate(images):
            files.append(('files', (f'image_{i}.png', img_data, 'image/png')))
        
        data = {
            'style': kwargs.get('style', 'anime'),
            'strength': kwargs.get('strength', 0.7),
            'guidance_scale': kwargs.get('guidance_scale', 7.5),
            'num_inference_steps': kwargs.get('num_inference_steps', 20)
        }
        
        try:
            response = self.session.post(
                f"{self.base_url}/colorize/batch",
                files=files,
                data=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                return result.get('batch_id')
            else:
                print(f"❌ Batch upload failed: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"❌ Batch upload error: {e}")
            return None
    
    def list_tasks(self, limit: int = 10) -> list:
        """List recent tasks"""
        try:
            response = self.session.get(
                f"{self.base_url}/colorize/tasks",
                params={'limit': limit},
                timeout=10
            )
            return response.json().get('tasks', []) if response.status_code == 200 else []
        except Exception as e:
            print(f"❌ List tasks error: {e}")
            return []


def demo_single_colorization():
    """Demo single image colorization"""
    print("🎨 Single Image Colorization Demo")
    print("=" * 50)
    
    client = MangaColorizerClient()
    
    # 1. Health check
    print("1. Checking server health...")
    if not client.health_check():
        print("❌ Server not available")
        return
    
    # 2. Create test image
    print("2. Creating test manga sketch...")
    image_data = client.create_test_image("manga")
    
    # 3. Upload and colorize
    print("3. Uploading and starting colorization...")
    task_id = client.upload_and_colorize(
        image_data,
        style="anime",
        strength=0.7,
        guidance_scale=7.5,
        num_inference_steps=15
    )
    
    if not task_id:
        return
    
    # 4. Wait for completion
    if client.wait_for_completion(task_id):
        # 5. Get result
        result = client.get_result(task_id)
        if result:
            print("✅ Colorization complete!")
            print(f"   Task ID: {task_id}")
            print(f"   Original: {result.get('original_image_url')}")
            print(f"   Colored: {result.get('colored_image_url')}")
            print(f"   Processing time: {result.get('processing_time', 0):.2f}s")
            print(f"   Style: {result.get('metadata', {}).get('style', 'unknown')}")


def demo_batch_processing():
    """Demo batch processing"""
    print("\n🎨 Batch Processing Demo")
    print("=" * 50)
    
    client = MangaColorizerClient()
    
    # Create multiple test images
    images = [
        client.create_test_image("manga"),
        client.create_test_image("landscape")
    ]
    
    print("1. Creating batch colorization...")
    batch_id = client.batch_colorize(
        images,
        style="cartoon",
        strength=0.8,
        guidance_scale=8.0,
        num_inference_steps=12
    )
    
    if batch_id:
        print(f"✅ Batch created: {batch_id}")
        
        # List recent tasks
        tasks = client.list_tasks(limit=5)
        print(f"   Recent tasks: {len(tasks)} found")


def demo_all_styles():
    """Demo all available styles"""
    print("\n🎨 All Styles Demo")
    print("=" * 50)
    
    client = MangaColorizerClient()
    
    styles = ["anime", "realistic", "cartoon", "sketch", "watercolor"]
    image_data = client.create_test_image("manga")
    
    for style in styles:
        print(f"Testing style: {style}")
        task_id = client.upload_and_colorize(
            image_data,
            style=style,
            strength=0.7,
            guidance_scale=7.5,
            num_inference_steps=10
        )
        
        if task_id and client.wait_for_completion(task_id, max_wait=30):
            result = client.get_result(task_id)
            if result:
                print(f"   ✅ {style}: {result.get('colored_image_url')}")


def main():
    """Run all demos"""
    print("🎨 Manga Colorizer - Complete Client Demo")
    print("=" * 60)
    
    # Check if server is running
    client = MangaColorizerClient()
    if not client.health_check():
        print("❌ Server not running. Please start the server first:")
        print("   python scripts/start_local.py")
        return
    
    # Run demos
    demo_single_colorization()
    demo_batch_processing()
    demo_all_styles()
    
    print("\n" + "=" * 60)
    print("✅ All demos completed successfully!")
    print("\n📚 Next steps:")
    print("   1. Use the API endpoints directly")
    print("   2. Integrate with mobile app")
    print("   3. Deploy to production")


if __name__ == "__main__":
    main()
